#pragma once

#include <opencv2\opencv.hpp>

using namespace cv;

class DTree
{
private:
	CvDTree dtree;

public:

	void set_trainer(CvDTree _dtree)
	{
		dtree = _dtree;
	}

	CvDTree get_trainer()
	{
		return dtree;
	}

	void train(const Mat& data, const Mat& label, const CvDTreeParams& params)
	{
		dtree.train(data, CV_ROW_SAMPLE, label, Mat(), Mat(), Mat(), Mat(), params);
	}

	float predict(const Mat& sample)
	{
		CvDTreeNode* resultNode = dtree.predict(sample);
		float resultValue = resultNode->value;

		return resultValue;
	}
};